import os
import cv2
import numpy as np
import codecs

image_dir = "/data/ppocr/trainData/test/images/"
gt_dir = "/data/ppocr/trainData/test/gt"
rec_base_dir = "/data/ppocr/trainData/test/rec"


def pick_top_left(poly):
    """
    bl  bottom left
    tl  top left
    br bottom right
    tr top right
    """
    idx = np.argsort(poly[:, 0])  # 升序
    if poly[idx[0], 1] < poly[idx[1], 1]:
        tl = idx[0]
        bl = idx[1]
    else:
        tl = idx[1]
        bl = idx[0]

    if poly[idx[-1], 1] < poly[idx[-2], 1]:
        tr = idx[-1]
        br = idx[-2]
    else:
        tr = idx[-2]
        br = idx[-1]
    return poly[(tl, tr, br, bl), :]


def get_rotate_crop_image(img, points):
    left = int(np.min(points[:, 0]))
    right = int(np.max(points[:, 0]))
    top = int(np.min(points[:, 1]))
    bottom = int(np.max(points[:, 1]))
    img_crop = img[top:bottom, left:right, :].copy()

    # cv2.imwrite("aa.jpg", img_crop)
    points[:, 0] = points[:, 0] - left
    points[:, 1] = points[:, 1] - top

    img_crop_width = int(np.linalg.norm(points[0] - points[1]))
    img_crop_height = int(np.linalg.norm(points[0] - points[3]))
    pts_std = np.float32([[0, 0], [img_crop_width, 0], [img_crop_width, img_crop_height], [0, img_crop_height]])
    M = cv2.getPerspectiveTransform(points, pts_std)
    dst_img = cv2.warpPerspective(
        img_crop,
        M, (img_crop_width, img_crop_height),
        borderMode=cv2.BORDER_REPLICATE)
    if img_crop_height * 1.0 / img_crop_width > 1:
        dst_img = np.rot90(dst_img)
    return dst_img


for name in os.listdir(image_dir):
    # name = "0000190.jpg"
    common_name = name.split(".")[0]
    gt_name = common_name + ".txt"
    image_path = os.path.join(image_dir, name)
    gt_path = os.path.join(gt_dir, gt_name)

    # H * W
    image = cv2.imread(image_path)
    lines = []

    begin_index = 0
    with codecs.open(gt_path, "r", "utf8") as f:
        for line in f.readlines():
            mark = [float(_i) for _i in line.strip().split(",")]
            points = [[mark[0], mark[1]], [mark[2], mark[3]], [mark[4], mark[5]], [mark[6], mark[7]]]
            # if begin_index == 13 or begin_index == 32:
            #     print(points)
            #     import pdb
            #     pdb.set_trace()
            # points = pick_top_left(np.array(points, dtype=np.float32))

            crop_image = get_rotate_crop_image(image, np.array(points, dtype=np.float32))
            begin_index += 1
            rec_dir = os.path.join(rec_base_dir, common_name)
            if not os.path.exists(rec_dir):
                os.makedirs(rec_dir)

            rec_path = os.path.join(rec_dir, str(begin_index).zfill(3))+".jpg"
            cv2.imwrite(rec_path, crop_image)
    # break